#include "camera.h"
#include <fstream>
#include <iostream>
#include <rapidjson/document.h>
#include "geometry.h"

cv::Mat Camera::get_K() const
{
    cv::Mat camera_K = (cv::Mat_<double>(3, 3) << K(0, 0), K(0, 1), K(0, 2),
                        K(1, 0), K(1, 1), K(1, 2),
                        K(2, 0), K(2, 1), K(2, 2));
    return camera_K;
}

cv::Mat Camera::get_D() const
{
    cv::Mat camera_D = (cv::Mat_<double>(5, 1) << D[0], D[1], D[2], D[3], D[4]);
    return camera_D;
}

bool Camera::LoadFromCfg(const std::string& cfg_file)
{
    std::ifstream ifs(cfg_file);
    if(!ifs.is_open())
    {
        std::cout << "open config file failed !" << std::endl;
        return false;
    }

    std::string json_content((std::istreambuf_iterator<char>(ifs)), std::istreambuf_iterator<char>());
    ifs.close();

    rapidjson::Document dom;
    dom.Parse<rapidjson::ParseFlag::kParseDefaultFlags>(json_content.c_str());
    if(dom.HasParseError())
    {
        std::cout << "parse config file failed !" << std::endl;
        return false;
    }

    width = dom["P"]["width"].GetDouble();
    height = dom["P"]["height"].GetDouble();
    
    double fx = dom["K"]["fx"].GetDouble();
    double fy = dom["K"]["fy"].GetDouble();
    double cx = dom["K"]["cx"].GetDouble();
    double cy = dom["K"]["cy"].GetDouble();

    K << fx, 0.0, cx,
         0.0, fy, cy,
         0.0, 0.0, 1.0;
    
    D[0] = dom["D"]["k1"].GetDouble();
    D[1] = dom["D"]["k2"].GetDouble();
    D[2] = dom["D"]["p1"].GetDouble();
    D[3] = dom["D"]["p2"].GetDouble();
    D[4] = dom["D"]["k3"].GetDouble();

    double x = dom["T"]["x"].GetDouble();
    double y = dom["T"]["y"].GetDouble();
    double z = dom["T"]["z"].GetDouble();
    mount.translation = Eigen::Vector3d(x, y, z);

    double roll = geometry::D2R(dom["R"]["roll"].GetDouble());
    double pitch = geometry::D2R(dom["R"]["pitch"].GetDouble());
    double yaw = geometry::D2R(dom["R"]["yaw"].GetDouble());
    mount.SetYpr(Eigen::Vector3d(roll, pitch, yaw));
    return true;
}

cv::Point2d Camera::RemoveDisto(double u, double v) const
{
	const double epsilon = 1e-8;
	cv::Vec2d p_u(u, v), p(u, v);
	cv::Vec2d d = DistoFunction(u, v);
	double err = cv::norm(p_u + d, p, 2);

	int count = 0;
	while (err > epsilon && count <= 50)
	{
		p_u = p - d;
		d = DistoFunction(p_u[0], p_u[1]);
		err = cv::norm(p_u + d, p, 2);
		count++;
	}
	return cv::Point2d(p_u[0], p_u[1]);
}

cv::Point2d Camera::RemoveDisto(const cv::Point2d & uv) const
{
	return RemoveDisto(uv.x, uv.y);
}

cv::Point2d Camera::AddDisto(double u, double v) const
{
    cv::Mat K = get_K(), D = get_D();
	cv::Mat xi3 = (cv::Mat_<double>(3, 1) << u, v, 1);
	cv::Mat xc3 = K.inv() * xi3;

	cv::Vec2d p3d;
	p3d[0] = xc3.at<double>(0, 0);
	p3d[1] = xc3.at<double>(1, 0);
	p3d = p3d + DistoFunction(p3d[0], p3d[1]);

	xc3 = (cv::Mat_<double>(3, 1) << p3d[0], p3d[1], 1);
	xi3 = K * xc3;
	xi3 / xi3.at<double>(2, 0);
	cv::Point2d p2d(xi3.at<double>(0, 0), xi3.at<double>(1, 0));
	return p2d;
}

cv::Vec2d Camera::DistoFunction(double u, double v) const
{
    cv::Mat D = get_D();
	const double k1 = D.at<double>(0, 0);
	const double k2 = D.at<double>(0, 1);
	const double k3 = D.at<double>(0, 4);
	const double p1 = D.at<double>(0, 2);
	const double p2 = D.at<double>(0, 3);

	const double x = u;
	const double y = v;

	const double r2 = x * x + y * y;
	const double r4 = r2 * r2;
	const double r6 = r4 * r2;
	const double k_diff = (k1 * r2 + k2 * r4 + k3 * r6);
	const double t_x = p2 * (r2 + 2 * x * x) + 2 * p1 * x * y;
	const double t_y = p1 * (r2 + 2 * y * y) + 2 * p2 * x * y;

	return (x * k_diff + t_x, y * k_diff + t_y);
}

cv::Point2d Camera::Pixel2Cam(const cv::Point2d & xi) const
{
    cv::Mat K = get_K();
	cv::Mat xi3  = (cv::Mat_<double>(3, 1) << xi.x, xi.y, 1);
	cv::Mat xc3 = K.inv() * xi3;

	cv::Point2d  xc;
	xc.x = xc3.at<double>(0, 0);
	xc.y = xc3.at<double>(1, 0);
	return xc;
}

cv::Point2d Camera::Cam2Pixel(const cv::Point2d & xc) const
{
    cv::Mat K = get_K();
	cv::Mat xc3 = (cv::Mat_<double>(3, 1) << xc.x, xc.y, 1);
	cv::Mat xi3 = K * xc3;

	cv::Point2d  xi;
	xi.x = xi3.at<double>(0, 0);
	xi.y = xi3.at<double>(1, 0);
	return xi;
}

bool Camera::Project3d_2d(Eigen::Vector3d pt3d, Eigen::Vector2d &img2D) const
{
	if (pt3d.z() > 0)
	{
		Eigen::Vector3d rpt = K * pt3d;
		rpt /= rpt.z();
		img2D = Eigen::Vector2d(rpt.x(), rpt.y());
		return true;
	}
	return false;
}